from funasr import AutoModel
import os

import soundfile as sf
import numpy as np

# 初始化模型
model = AutoModel(model="fsmn-vad")


def process_multiple_files(file_list, chunk_size=50):
    """
    将多个音频文件合并后处理
    :param file_list: 输入的音频文件列表
    :param chunk_size: 分块大小（毫秒）
    """
    combined_audio = []
    sample_rate = None

    # 加载所有音频文件
    for file in file_list:
        audio, sr = sf.read(file)
        if sample_rate is None:
            sample_rate = sr
        elif sample_rate != sr:
            raise ValueError("所有音频文件的采样率必须一致")
        combined_audio.append(audio)

    # 合并所有音频
    combined_audio = np.concatenate(combined_audio)

    # 开始处理
    chunk_stride = int(chunk_size * sample_rate / 1000)
    cache = {}
    total_chunk_num = int(len(combined_audio) / chunk_stride + 1)

    for i in range(total_chunk_num):
        speech_chunk = combined_audio[i * chunk_stride:(i + 1) * chunk_stride]
        is_final = i == total_chunk_num - 1
        res = model.generate(input=speech_chunk, cache=cache, is_final=is_final, chunk_size=chunk_size)

        if len(res[0]["value"]):
            if is_end(res):
                print('isend at ' + str(i))
                print(res)
            else:
                print('not end')
                print(res)

def is_end(data):
    # 遍历列表中的每个字典
    for item in data:
        # 检查字典中是否有 'value' 键
        if 'value' in item:
            # 获取 'value' 的值
            value = item['value']
            # 确保 'value' 是一个列表，并且至少有一个子列表
            if isinstance(value, list) and len(value) > 0:
                # 遍历 'value' 中的每个子列表
                for sublist in value:
                    # 检查子列表是否是列表，并且长度为2
                    if isinstance(sublist, list) and len(sublist) == 2:
                        # 检查第一个元素是否为 -1，第二个元素是否大于 0
                        if sublist[0] == -1 and sublist[1] > 0:
                            return True
    return False

if __name__ == '__main__':
    # directory = "/home/leo/desktop/test/ll/var/input/11/"
    directory = "/data/kootxt/local-logs/var/input/f83ef92d8f4544969a35bd40d0bf0bb1/"
    files = [f for f in os.listdir(directory) if os.path.isfile(os.path.join(directory, f))]
    sorted_files = sorted(files, key=lambda x: int(x))
    # 用于存储字节数组的列表
    byte_arrays = []

    for file_name in sorted_files:
        file_path = os.path.join(directory, file_name)
        if os.path.isfile(file_path):
            byte_arrays.append(file_path)

    process_multiple_files(byte_arrays)